import asyncio
import os
import subprocess
from playwright.async_api import async_playwright
from utils import load_config, get_md5, ensure_dir
import logging
import time

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class WebCrawler:
    def __init__(self, config):
        self.config = config
        self.browser = None
        self.chrome_process = None
        ensure_dir(config['save']['output_dir'])

    def start_chrome(self):
        """启动Chrome浏览器"""
        try:
            # 获取当前用户的Chrome数据目录
            default_user_data_dir = os.path.expandvars(r'%LOCALAPPDATA%\Google\Chrome\User Data')
            
            cmd = [
                self.config['browser']['chrome_path'],
                f"--remote-debugging-port=9222",
                f"--user-data-dir={default_user_data_dir}",  # 使用默认的用户数据目录
                f"--profile-directory={self.config['browser']['profile']}",
                "--no-first-run",
                "--no-default-browser-check"
            ]
            logging.info(f"正在启动Chrome浏览器，使用配置目录: {default_user_data_dir}")
            
            # 检查是否有其他Chrome实例在运行
            self.kill_existing_chrome()
            
            self.chrome_process = subprocess.Popen(cmd)
            time.sleep(3)
            return True
        except Exception as e:
            logging.error(f"启动Chrome失败: {str(e)}")
            return False

    def kill_existing_chrome(self):
        """关闭所有正在运行的Chrome进程"""
        try:
            if os.name == 'nt':  # Windows
                os.system('taskkill /F /IM chrome.exe /T')
            else:  # Linux/Mac
                os.system('pkill -f chrome')
            time.sleep(2)  # 等待进程完全关闭
        except Exception as e:
            logging.warning(f"关闭现有Chrome进程失败: {str(e)}")

    async def init_browser(self):
        try:
            # 获取当前用户的Chrome数据目录
            default_user_data_dir = os.path.expandvars(r'%LOCALAPPDATA%\Google\Chrome\User Data')
            
            logging.info(f"正在启动Chrome浏览器，使用配置目录: {default_user_data_dir}")
            
            # 检查是否有其他Chrome实例在运行
            self.kill_existing_chrome()
            
            # 使用playwright直接启动浏览器
            playwright = await async_playwright().start()
            self.browser = await playwright.chromium.launch_persistent_context(
                user_data_dir=default_user_data_dir,
                channel="chrome",
                headless=False,
                args=[
                    '--remote-debugging-port=9222',
                    f'--profile-directory={self.config["browser"]["profile"]}',
                    '--no-first-run',
                    '--no-default-browser-check',
                    '--password-store=basic'
                ]
            )
            
            logging.info("成功启动Chrome浏览器")
            
        except Exception as e:
            logging.error(f"浏览器启动失败: {str(e)}")
            raise

    async def save_page(self, url):
        try:
            # 使用已有的上下文创建新页面
            page = await self.browser.new_page()
            logging.info(f"正在访问页面: {url}")
            
            await page.goto(url, wait_until='networkidle', timeout=60000)
            await page.wait_for_load_state('networkidle')
            
            html_content = await page.content()
            
            filename = f"{get_md5(url)}.html"
            filepath = os.path.join(self.config['save']['output_dir'], filename)
            
            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(html_content)
                
            logging.info(f"成功保存页面到: {filepath}")
            
        except Exception as e:
            logging.error(f"处理页面失败 {url}: {str(e)}")
        finally:
            if page:
                await page.close()

    async def process_urls(self):
        try:
            await self.init_browser()
            tasks = [self.save_page(url) for url in self.config['urls']]
            await asyncio.gather(*tasks)
        finally:
            if self.browser:
                await self.browser.close()
                logging.info("已关闭Chrome浏览器")

async def main():
    config = load_config('config.yaml')
    crawler = WebCrawler(config)
    await crawler.process_urls()

if __name__ == "__main__":
    asyncio.run(main()) 